import pickle
import os

trainCmd="java -Xmx4g -jar RankLib-2.1-patched.jar -ranker 2 -metric2t NDCG@100 -train train.txt -tvs 0.8 -save model.txt > train_preci.txt"
validCmd="java -Xmx4g -jar RankLib-2.1-patched.jar -ranker 2 -metric2t NDCG@100 -train train.txt -tvs 0.8 -test test.txt > valid_preci%d.txt"
#trainCmd="java -Xmx4g -jar RankLib-2.1-patched.jar -ranker 6 -train train.txt -metric2t NDCG@100 -save model.txt -tvs 0.8 >> train_preci.txt"
rankCmd="java -Xmx4g -jar RankLib-2.1-patched.jar -load model.txt -rank test.txt -score output.txt"

# The line below, only for debug
#fd=pickle.load(open("feature_dict","r"))

# fds is a vector of feature_dicts
# Assume the first element of each feature vector is the score
# So the first item in training vector is 5,1,0
# while for testing feature vector, the first item is assumed 0
def FDsToFile(fds,fname):
    fout=open(fname,"w")
    for k in range(0,len(fds)):
        for id in fds[k]:
            fd=fds[k]
            fout.write(str(fd[id][0])+" qid:"+str(k))
            for i in range(1,len(fd[id])):
                fout.write(" "+str(i)+":"+str(fd[id][i]))
            fout.write("\n")
    fout.close()

def trainModel(fds, validate=False):
    if validate:
        for k in range(0,len(fds)):
            # validate k-th list using the others to train
            FDsToFile([fds[k]],"test.txt")
            FDsToFile(fds[0:k]+fds[k+1:],"train.txt")
            os.system(validCmd % k)
    FDsToFile(fds,"train.txt")
    os.system(trainCmd);

def loadOutput(fd):
    out=open("output.txt","r")
    res=[]
    for id in fd:
        line=out.readline().split("\t")
        res+=[(id, float(line[2]))]
    res.sort(key=lambda x:x[1], reverse=True)
    return map(lambda x:x[0], res)

def testOne(fd):
    FDsToFile([fd],"test.txt")
    os.system(rankCmd)
    return loadOutput(fd)
